import os
import pandas as pd
from langchain.prompts import PromptTemplate
from langchain.schema import BaseOutputParser
from langchain.chains import LLMChain
from langchain_openai import ChatOpenAI
from typing import List, Dict
import json
import re
import glob
import shutil
from dotenv import load_dotenv

load_dotenv()

# 阿里云百炼配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

# 验证API密钥
if not DASHSCOPE_CONFIG["api_key"]:
    raise ValueError("请设置DASHSCOPE_API_KEY环境变量")


class QAOutputParser(BaseOutputParser):
    """解析LLM输出的问答对"""

    def parse(self, text: str) -> List[Dict[str, str]]:
        """解析问答对文本"""
        qa_pairs = []

        # 预处理：提取JSON内容（处理代码块格式）
        json_text = self._extract_json_content(text)

        # 尝试解析JSON格式
        try:
            if json_text.strip().startswith('[') or json_text.strip().startswith('{'):
                data = json.loads(json_text)
                if isinstance(data, list):
                    return data
                elif isinstance(data, dict) and 'qa_pairs' in data:
                    return data['qa_pairs']
        except json.JSONDecodeError:
            pass

        # 解析文本格式的问答对
        lines = text.strip().split('\n')
        current_qa = {}

        for line in lines:
            line = line.strip()
            if not line:
                continue

            # 匹配问题格式 - 提取纯净的问题内容
            q_match = re.match(r'^(?:问题?\d*|Q\d*|Question\d*)\s*[：:：]?\s*(.+)$', line, re.IGNORECASE)
            if q_match:
                if current_qa and 'question' in current_qa and 'answer' in current_qa:
                    qa_pairs.append(current_qa)
                # 清理问题内容，移除可能的标号和冒号
                clean_question = self._clean_content(q_match.group(1))
                current_qa = {'question': clean_question}
                continue

            # 匹配答案格式 - 提取纯净的答案内容
            a_match = re.match(r'^(?:答案?\d*|A\d*|Answer\d*)\s*[：:：]?\s*(.+)$', line, re.IGNORECASE)
            if a_match and 'question' in current_qa:
                # 清理答案内容，移除可能的标号和冒号
                clean_answer = self._clean_content(a_match.group(1))
                current_qa['answer'] = clean_answer
                continue

            # 如果当前有问题但没有答案，将此行作为答案
            if 'question' in current_qa and 'answer' not in current_qa:
                clean_answer = self._clean_content(line)
                current_qa['answer'] = clean_answer

        # 添加最后一个问答对
        if current_qa and 'question' in current_qa and 'answer' in current_qa:
            qa_pairs.append(current_qa)

        return qa_pairs

    def _extract_json_content(self, text: str) -> str:
        """提取JSON内容，处理代码块格式"""
        text = text.strip()

        # 检查是否包含```json代码块
        json_block_pattern = r'```json\s*\n?([\s\S]*?)\n?```'
        json_match = re.search(json_block_pattern, text, re.IGNORECASE)

        if json_match:
            # 提取代码块中的JSON内容
            return json_match.group(1).strip()

        # 如果没有代码块，返回原文本
        return text

    def _clean_content(self, content: str) -> str:
        """清理内容，移除标号和多余的冒号"""
        # 移除开头的数字标号和冒号
        content = re.sub(r'^\d+[.、：:：]?\s*', '', content)
        # 移除开头的字母标号和冒号
        content = re.sub(r'^[A-Za-z][.、：:：]?\s*', '', content)
        # 移除多余的冒号
        content = re.sub(r'^[：:：]+\s*', '', content)
        return content.strip()


class JSONQAGenerator:
    """JSON问答对生成器"""

    def __init__(self):
        # 初始化阿里云百炼模型
        self.llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.7
        )

        # 创建自动生成数量的提示模板
        self.auto_prompt_template = PromptTemplate(
            input_variables=["text"],
            template="""你是一个专业的问答生成专家。请根据以下文本内容，全面地生成中文问答对。

要求：
1. 问答对必须全面覆盖文本的所有关键信息点，不遗漏重要内容，问题要有意义
2. 问答对必须使用原话，不能添加或修改信息
3. 根据文本内容信息丰富程度，自主决定生成问答对的数量
4. 如果分段或内部标题能够完整总结分段内容，增加以标题生成问题，分段内容作为回答的问答对，注意标题不带章节数字号

文本内容：
{text}

请以标准JSON格式输出问答对：

{{
  "qa_pairs": [
    {{
      "question": " ",
      "answer": " "
    }},
    {{
      "question": " ",
      "answer": " "
    }},
    {{
      "question": " ",
      "answer": " "
    }}
  ]
}}

注意：
1. 只输出JSON格式的结果
2. 确保JSON格式正确，使用双引号
3. 问题和答案内容中的引号需要转义
4. 不要输出其他格式的内容
"""
        )

        # 创建自动生成LLM链
        self.auto_qa_chain = LLMChain(
            llm=self.llm,
            prompt=self.auto_prompt_template,
            output_parser=QAOutputParser()
        )

    def auto_generate_qa_pairs(self, text: str) -> List[Dict[str, str]]:
        """自动生成问答对，数量由模型根据文本内容决定"""
        try:
            result = self.auto_qa_chain.run(text=text)
            return result if isinstance(result, list) else []
        except Exception as e:
            print(f"自动生成问答对时出错: {e}")
            return []

    def process_json_file(self, json_file_path: str, output_file: str):
        """处理单个JSON文件"""
        print(f"正在处理文件: {os.path.basename(json_file_path)}")
        
        try:
            # 读取JSON文件
            with open(json_file_path, 'r', encoding='utf-8') as f:
                json_data = json.load(f)
            
            # 检查JSON文件格式
            if 'document_sections' not in json_data:
                print(f"文件 {os.path.basename(json_file_path)} 格式不正确，缺少 'document_sections' 字段")
                return
            
            sections = json_data['document_sections']
            category_id = os.path.splitext(os.path.basename(json_file_path))[0]
            
            print(f"找到 {len(sections)} 个分段")
            
            all_qa_pairs = []
            
            # 处理每个分段
            for i, section in enumerate(sections, 1):
                section_id = section.get('section_id', '未知分段')
                # 将section_id转换为字符串再进行切片
                section_display = str(section_id)[:30] if section_id != '未知分段' else section_id
                print(f"  处理分段 {i}/{len(sections)}: {section_display}...")
                
                # 构建分段文本
                h1_title = section.get('h1_title', '') or ''
                h2_title = section.get('h2_title', '') or ''
                content = section.get('content', '') or ''
                
                if not content.strip():
                    print(f"    分段 {i} 内容为空，跳过")
                    continue
                
                # 组合标题和内容
                section_text = ""
                if h1_title:
                    section_text += f"一级标题：{h1_title}\n"
                if h2_title:
                    section_text += f"二级标题：{h2_title}\n"
                section_text += f"内容：{content}"
                
                # 生成问答对
                qa_pairs = self.auto_generate_qa_pairs(section_text)
                
                if qa_pairs:
                    # 为每个问答对添加分段信息
                    for qa_index, qa in enumerate(qa_pairs, 1):
                        qa['section_id'] = section.get('section_id', i)
                        qa['h1_title'] = h1_title
                        qa['h2_title'] = h2_title
                        qa['qa_index_in_section'] = qa_index
                        
                        # 添加context字段 - 只包含当前分段的信息
                        current_section_context = {
                            'section_id': section.get('section_id', i),
                            'h1_title': h1_title,
                            'h2_title': h2_title,
                            'content': content
                        }
                        
                        qa['context'] = json.dumps(current_section_context, ensure_ascii=False)
                    
                    all_qa_pairs.extend(qa_pairs)
                    print(f"    成功生成 {len(qa_pairs)} 个问答对")
                else:
                    print(f"    分段 {i} 未能生成问答对")
            
            # 保存到Excel文件
            if all_qa_pairs:
                self.save_to_excel(all_qa_pairs, output_file, category_id)
                print(f"文件 {os.path.basename(json_file_path)} 处理完成，共生成 {len(all_qa_pairs)} 个问答对")
            else:
                print(f"文件 {os.path.basename(json_file_path)} 未生成任何问答对")
                
        except Exception as e:
            print(f"处理文件 {os.path.basename(json_file_path)} 时出错: {e}")

    def process_json_directory(self, input_dir: str, output_dir: str):
        """批量处理目录中的所有JSON文件"""
        # 查找所有json文件
        json_pattern = os.path.join(input_dir, "*.json")
        json_files = glob.glob(json_pattern)
        
        if not json_files:
            print(f"在目录 {input_dir} 中未找到任何.json文件")
            return
        
        print(f"找到 {len(json_files)} 个JSON文件")
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        total_qa_pairs = 0
        
        for i, json_file in enumerate(json_files, 1):
            print(f"\n正在处理第 {i}/{len(json_files)} 个文件: {os.path.basename(json_file)}")
            
            # 生成输出文件名
            base_name = os.path.splitext(os.path.basename(json_file))[0]
            output_file = os.path.join(output_dir, f"{base_name}.xlsx")
            
            # 如果输出文件已存在，先删除
            if os.path.exists(output_file):
                os.remove(output_file)
            
            try:
                # 处理JSON文件
                self.process_json_file(json_file, output_file)
                
                # 统计生成的问答对数量
                if os.path.exists(output_file):
                    df = pd.read_excel(output_file)
                    file_qa_count = len(df)
                    total_qa_pairs += file_qa_count
                    print(f"文件 {base_name}.xlsx 已保存，包含 {file_qa_count} 个问答对")
                
            except Exception as e:
                print(f"处理文件 {os.path.basename(json_file)} 时出错: {e}")
                continue
        
        print(f"\n批量处理完成！")
        print(f"总共处理了 {len(json_files)} 个JSON文件")
        print(f"总共生成了 {total_qa_pairs} 个问答对")
        print(f"结果已保存到目录: {output_dir}")

    def save_to_excel(self, qa_pairs: List[Dict[str, str]], filename: str, category_id: str = None):
        """保存问答对到Excel文件"""
        if not qa_pairs:
            print("没有问答对可保存")
            return
        
        # 准备数据
        data = []
        for i, qa in enumerate(qa_pairs, 1):
            row = {
                '标号': i,
                '提问': qa.get('question', ''),
                '回答': qa.get('answer', '')
            }
            if category_id:
                row['category_id'] = category_id
            # 添加分段信息
            if 'section_id' in qa:
                row['分段ID'] = qa['section_id']
            if 'h1_title' in qa:
                row['一级标题'] = qa['h1_title']
            if 'h2_title' in qa:
                row['二级标题'] = qa['h2_title']
            if 'qa_index_in_section' in qa:
                row['分段内序号'] = qa['qa_index_in_section']
            # 添加context字段
            if 'context' in qa:
                row['context'] = qa['context']
            data.append(row)
        
        # 创建DataFrame
        columns = ['标号', '提问', '回答']
        if category_id:
            columns.append('category_id')
        if any('分段ID' in row for row in data):
            columns.append('分段ID')
        if any('一级标题' in row for row in data):
            columns.append('一级标题')
        if any('二级标题' in row for row in data):
            columns.append('二级标题')
        if any('分段内序号' in row for row in data):
            columns.append('分段内序号')
        if any('context' in row for row in data):
            columns.append('context')
        
        df = pd.DataFrame(data, columns=columns)
        df.to_excel(filename, index=False)


def clear_folder(folder_path: str):
    """清空指定文件夹中的所有文件"""
    if os.path.exists(folder_path):
        try:
            shutil.rmtree(folder_path)
            os.makedirs(folder_path, exist_ok=True)
            print(f"已清空文件夹: {folder_path}")
        except Exception as e:
            print(f"清空文件夹失败 {folder_path}: {e}")
    else:
        os.makedirs(folder_path, exist_ok=True)
        print(f"已创建文件夹: {folder_path}")

def main():
    """主函数 - 批量处理JSON文件"""
    generator = JSONQAGenerator()
    
    print("=== JSON问答对生成工具 ===")
    print(f"连接到阿里云百炼服务: {DASHSCOPE_CONFIG['base_url']}")
    print(f"使用模型: {DASHSCOPE_CONFIG['model']}")
    print()
    
    # 设置输入目录和输出目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    input_dir = os.path.join(os.path.dirname(current_dir), "A_segmenter", "json")
    output_dir = os.path.join(current_dir, "jsonQA")
    
    print(f"当前工作目录: {current_dir}")
    print(f"输入目录: {input_dir}")
    print(f"输出目录: {output_dir}")
    print()
    
    # 检查输入目录是否存在
    if not os.path.exists(input_dir):
        print(f"错误：输入目录不存在 - {input_dir}")
        return
    
    # 清空输出目录
    clear_folder(output_dir)
    
    try:
        # 批量处理目录中的所有JSON文件
        generator.process_json_directory(input_dir, output_dir)
        
    except Exception as e:
        print(f"批量处理过程中发生错误: {e}")


if __name__ == "__main__":
    main()